#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
MongoDB数据导入脚本
将民法典数据导入到MongoDB数据库
"""

import json
import os
from pymongo import MongoClient
from datetime import datetime
from typing import List, Dict, Any

class MongoDBImporter:
    """MongoDB数据导入器"""
    
    def __init__(self, connection_string: str = "mongodb://localhost:27017/", db_name: str = "civil_code_db"):
        self.connection_string = connection_string
        self.db_name = db_name
        self.client = None
        self.db = None
        self.collection = None
    
    def connect(self):
        """连接MongoDB"""
        try:
            self.client = MongoClient(self.connection_string)
            self.db = self.client[self.db_name]
            self.collection = self.db.civil_code_articles
            
            # 测试连接
            self.client.admin.command('ping')
            print(f"✅ 成功连接到MongoDB: {self.db_name}")
            return True
            
        except Exception as e:
            print(f"❌ MongoDB连接失败: {str(e)}")
            print("💡 请确保MongoDB服务已启动")
            print("   - Windows: 启动MongoDB服务")
            print("   - Linux/Mac: sudo systemctl start mongod 或 brew services start mongodb")
            return False
    
    def create_collections_and_indexes(self):
        """创建集合和索引"""
        print("🔍 创建数据库集合和索引...")
        
        try:
            # 创建民法典条文集合
            articles_collection = self.db.civil_code_articles
            
            # 创建用户收藏集合
            favorites_collection = self.db.user_favorites
            
            # 创建搜索历史集合
            history_collection = self.db.search_history
            
            # 为条文集合创建索引
            articles_collection.create_index([
                ("content", "text"),
                ("title", "text"), 
                ("explanation", "text"),
                ("keywords", "text")
            ], name="fulltext_search")
            
            articles_collection.create_index("article_number", unique=True)
            articles_collection.create_index([
                ("chapter_info.part", 1),
                ("chapter_info.chapter_number", 1)
            ])
            articles_collection.create_index("keywords")
            
            # 为收藏集合创建索引
            favorites_collection.create_index([("user_id", 1), ("article_id", 1)], unique=True)
            favorites_collection.create_index("user_id")
            favorites_collection.create_index("created_at")
            
            # 为搜索历史集合创建索引
            history_collection.create_index("user_id")
            history_collection.create_index("created_at")
            history_collection.create_index("query")
            
            print("✅ 集合和索引创建完成")
            return True
            
        except Exception as e:
            print(f"❌ 创建索引失败: {str(e)}")
            return False
    
    def import_from_json(self, json_file_path: str):
        """从JSON文件导入数据"""
        if not os.path.exists(json_file_path):
            print(f"❌ JSON文件不存在: {json_file_path}")
            return False
        
        try:
            with open(json_file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            articles = data.get('articles', [])
            if not articles:
                print("❌ JSON文件中没有找到articles数据")
                return False
            
            return self.import_articles(articles)
            
        except Exception as e:
            print(f"❌ 读取JSON文件失败: {str(e)}")
            return False
    
    def import_articles(self, articles: List[Dict[str, Any]]):
        """导入条文数据"""
        if not articles:
            print("❌ 没有数据可导入")
            return False
        
        try:
            # 清空现有数据
            result = self.collection.delete_many({})
            print(f"🗑️ 已删除现有数据: {result.deleted_count} 条")
            
            # 准备数据
            for article in articles:
                # 移除_id字段，让MongoDB自动生成
                if '_id' in article:
                    del article['_id']
                if 'created_at' not in article:
                    article['created_at'] = datetime.now()
                if 'updated_at' not in article:
                    article['updated_at'] = datetime.now()
            
            # 批量插入数据
            result = self.collection.insert_many(articles)
            print(f"✅ 成功导入 {len(result.inserted_ids)} 条法条")
            
            # 验证导入结果
            total_count = self.collection.count_documents({})
            print(f"📊 数据库中共有 {total_count} 条法条")
            
            # 显示一些统计信息
            self.show_statistics()
            
            return True
            
        except Exception as e:
            print(f"❌ 数据导入失败: {str(e)}")
            return False
    
    def show_statistics(self):
        """显示数据统计信息"""
        print("\n📊 数据库统计信息:")
        
        try:
            # 按编统计
            pipeline = [
                {"$group": {
                    "_id": "$chapter_info.part",
                    "count": {"$sum": 1}
                }},
                {"$sort": {"_id": 1}}
            ]
            
            part_stats = list(self.collection.aggregate(pipeline))
            for stat in part_stats:
                print(f"   - {stat['_id']}: {stat['count']} 条")
            
            # 关键词统计
            total_keywords = self.collection.aggregate([
                {"$unwind": "$keywords"},
                {"$group": {"_id": None, "total": {"$sum": 1}}}
            ])
            
            keyword_result = list(total_keywords)
            if keyword_result:
                print(f"   - 总关键词数: {keyword_result[0]['total']} 个")
            
        except Exception as e:
            print(f"⚠️ 统计信息获取失败: {str(e)}")
    
    def test_queries(self):
        """测试查询功能"""
        print("\n🔍 测试数据库查询功能:")
        
        try:
            # 测试1: 按条文号查询
            article = self.collection.find_one({"article_number": "第一条"})
            if article:
                print(f"✅ 条文号查询: 找到 {article['article_number']}")
            else:
                print("❌ 条文号查询失败")
            
            # 测试2: 全文搜索
            search_results = list(self.collection.find(
                {"$text": {"$search": "民事主体"}},
                {"article_number": 1, "title": 1}
            ).limit(3))
            
            if search_results:
                print(f"✅ 全文搜索: 找到 {len(search_results)} 条相关法条")
                for result in search_results:
                    print(f"   - {result['article_number']}: {result['title'][:50]}...")
            else:
                print("❌ 全文搜索失败")
            
            # 测试3: 按章节查询
            chapter_results = list(self.collection.find(
                {"chapter_info.part": "第一编 总则"},
                {"article_number": 1}
            ).limit(5))
            
            if chapter_results:
                print(f"✅ 章节查询: 第一编总则有 {len(chapter_results)} 条法条（显示前5条）")
            else:
                print("❌ 章节查询失败")
                
        except Exception as e:
            print(f"❌ 查询测试失败: {str(e)}")
    
    def close(self):
        """关闭数据库连接"""
        if self.client:
            self.client.close()
            print("📪 MongoDB连接已关闭")

def main():
    """主函数"""
    print("=" * 60)
    print("💾 MongoDB数据导入工具")
    print("=" * 60)
    
    # 创建导入器
    importer = MongoDBImporter()
    
    # 连接数据库
    if not importer.connect():
        return
    
    try:
        # 创建集合和索引
        importer.create_collections_and_indexes()
        
        # 查找数据文件
        data_files = [
            "data/civil_code_full_processed.json",
            "data/civil_code_processed.json",
            "data/civil_code_crawled_processed.json"
        ]
        
        imported = False
        for data_file in data_files:
            if os.path.exists(data_file):
                print(f"\n📂 找到数据文件: {data_file}")
                if importer.import_from_json(data_file):
                    imported = True
                    break
        
        if not imported:
            print("❌ 没有找到可用的数据文件")
            print("💡 请先运行数据处理脚本生成数据文件")
            return
        
        # 测试查询功能
        importer.test_queries()
        
        print("\n🎉 MongoDB数据导入完成！")
        print("✅ 民法典数据已成功保存到数据库")
        
    finally:
        importer.close()

if __name__ == "__main__":
    main()
